package com.sfang.common.base.thread;

import cn.hutool.core.collection.CollUtil;
import cn.hutool.core.thread.NamedThreadFactory;
import org.slf4j.MDC;
import org.springframework.web.context.request.RequestAttributes;
import org.springframework.web.context.request.RequestContextHolder;

import java.util.Map;
import java.util.Objects;
import java.util.concurrent.*;

/**
 * 包装线程池
 * <p>
 * 提供简单的线程池创建方法
 * 提供trackId透传功能
 * </p>
 *
 * @author xuzilou
 * @date 2021/4/7 9:52
 */
public class MyThreadPool extends ThreadPoolExecutor {

    private MyThreadPool(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue);
    }

    public MyThreadPool(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory);
    }

    public MyThreadPool(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, handler);
    }

    public MyThreadPool(int corePoolSize, int maximumPoolSize, long keepAliveTime, TimeUnit unit, BlockingQueue<Runnable> workQueue, ThreadFactory threadFactory, RejectedExecutionHandler handler) {
        super(corePoolSize, maximumPoolSize, keepAliveTime, unit, workQueue, threadFactory, handler);
    }

    /**
     * @param corePoolSize 核心线程池大小
     * @return {@link ThreadPoolExecutor}
     */
    public static MyThreadPool createPool(int corePoolSize) {
        return new MyThreadPool(corePoolSize, corePoolSize, 1, TimeUnit.MINUTES
                , new LinkedBlockingQueue<>(), new NamedThreadFactory("myThreadPool-", false));
    }

    /**
     * @param corePoolSize 核心线程池大小
     * @param factory      线程生成工程
     * @return {@link ThreadPoolExecutor}
     */
    public static MyThreadPool createPool(int corePoolSize, ThreadFactory factory) {
        return new MyThreadPool(corePoolSize, corePoolSize, 1, TimeUnit.MINUTES
                , new LinkedBlockingQueue<>(), factory);
    }

    /**
     * @param corePoolSize 核心线程池大小
     * @param factory      线程生成工程
     * @param queue        队列数量
     * @return {@link ThreadPoolExecutor}
     */
    public static MyThreadPool createPool(int corePoolSize, int queue, ThreadFactory factory) {
        return new MyThreadPool(corePoolSize, corePoolSize, 1, TimeUnit.MINUTES
                , new LinkedBlockingQueue<>(queue), factory);
    }

    /**
     * @param corePoolSize             核心线程池大小
     * @param factory                  线程生成工程
     * @param blockingQueue            阻塞队列
     * @param rejectedExecutionHandler 拒绝策略
     * @return {@link ThreadPoolExecutor}
     */
    public static MyThreadPool createPool(int corePoolSize, BlockingQueue<Runnable> blockingQueue, ThreadFactory factory, RejectedExecutionHandler rejectedExecutionHandler) {
        return new MyThreadPool(corePoolSize, corePoolSize, 1, TimeUnit.MINUTES
                , blockingQueue, factory, rejectedExecutionHandler);
    }

    /**
     * 创建CPU密集型任务池
     * <p>
     * 根据公式:CPU密集型任务线程数=cpu核心数+1
     * 建议:计算型任务使用该方法创建线程池
     * </p>
     *
     * @return {@link ThreadPoolExecutor}
     */
    public static MyThreadPool createCpuDensePool() {
        int corePoolSize = Runtime.getRuntime().availableProcessors() + 1;
        return new MyThreadPool(corePoolSize, corePoolSize, 1, TimeUnit.MINUTES
                , new LinkedBlockingQueue<>(), new NamedThreadFactory("myThreadPool-", false));
    }

    /**
     * 创建IO密集型任务池
     * <p>
     * 根据公式:io密集型任务线程数=cpu核心数*2+1
     * 建议:数据库连接池或者远程调用使用该方法创建线程池
     * 实际应用场景建议进行压测通过jprofiler观察cpu耗时和cpu等待时间通过公式 (cpu等待时间/cpu耗时+1)*cpu核心数
     * 并且编写线程池流控,提醒并且提供手动设置线程池方法
     * </p>
     *
     * @return {@link ThreadPoolExecutor}
     */
    public static MyThreadPool createIoDensePool() {
        int corePoolSize = Runtime.getRuntime().availableProcessors() * 2 + 1;
        return new MyThreadPool(corePoolSize, corePoolSize, 1, TimeUnit.MINUTES
                , new LinkedBlockingQueue<>(), new NamedThreadFactory("myThreadPool-", false));
    }

    @Override
    public void execute(Runnable runnable) {
        super.execute(warp(runnable));
    }

    @Override
    protected void beforeExecute(Thread t, Runnable r) {
    }

    @Override
    protected void afterExecute(Runnable r, Throwable t) {
    }

    /**
     * 包装Runnable
     *
     * @param run Runnable
     * @return {@link Runnable}
     */
    public Runnable warp(Runnable run) {
        RequestAttributes requestAttributes = RequestContextHolder.getRequestAttributes();
        Map<String, String> copyOfContextMap = MDC.getCopyOfContextMap();
        return () -> {
            try {
                if (CollUtil.isNotEmpty(copyOfContextMap)) {
                    MDC.setContextMap(copyOfContextMap);
                }
                if (Objects.nonNull(requestAttributes)) {
                    RequestContextHolder.setRequestAttributes(requestAttributes);
                }
                run.run();
            } finally {
                MDC.clear();
            }
        };
    }
}
